import configargparse


def get_opts_base():
    parser = configargparse.ArgParser(config_file_parser_class=configargparse.YAMLConfigFileParser)
    parser.add_argument('--config_file', is_config_file=True)

    parser.add_argument('--dataset_type', type=str, default='filesystem', choices=['filesystem', 'memory'],
                        help="""specifies whether to hold all images in CPU memory during training, or whether to write randomized
                        batches or pixels/rays to disk""")
    parser.add_argument('--chunk_paths', type=str, nargs='+', default=None,
                        help="""scratch directory to write shuffled batches to when training using the filesystem dataset. 
    Should be set to a non-existent path when first created, and can then be reused by subsequent training runs once all chunks are written""")
    parser.add_argument('--num_chunks', type=int, default=200,
                        help='number of shuffled chunk files to write to disk. Each chunk should be small enough to fit into CPU memory')
    parser.add_argument('--disk_flush_size', type=int, default=10000000)
    parser.add_argument('--train_every', type=int, default=1,
                        help='if set to larger than 1, subsamples each n training images')

    parser.add_argument('--cluster_mask_path', type=str, default=None,
                        help='directory containing pixel masks for all training images (generated by create_cluster_masks.py)')

    parser.add_argument('--ckpt_path', type=str, default=None, help='path towards serialized model checkpoint')
    parser.add_argument('--container_path', type=str, default=None,
                        help='path towards merged Mega-NeRF model generated by merged_submodules.py')

    parser.add_argument('--near', type=float, default=1, help='ray near bounds')
    parser.add_argument('--far', type=float, default=None,
                        help='ray far bounds. Will be automatically set if not explicitly set')
    parser.add_argument('--ray_altitude_range', nargs='+', type=float, default=None,
                        help='constrains ray sampling to the given altitude')
    parser.add_argument('--coarse_samples', type=int, default=256,
                        help='number of coarse samples')
    parser.add_argument('--fine_samples', type=int, default=512,
                        help='number of additional fine samples')

    parser.add_argument('--train_scale_factor', type=int, default=1,
                        help='downsamples training images if greater than 1')
    parser.add_argument('--val_scale_factor', type=int, default=4,
                        help='downsamples validation images if greater than 1')

    parser.add_argument('--pos_xyz_dim', type=int, default=12,
                        help='frequency encoding dimension applied to xyz position')
    parser.add_argument('--pos_dir_dim', type=int, default=4,
                        help='frequency encoding dimension applied to view direction (set to 0 to disable)')
    parser.add_argument('--layers', type=int, default=8, help='number of layers in MLP')
    parser.add_argument('--skip_layers', type=int, nargs='+', default=[4], help='indices of the skip connections')
    parser.add_argument('--layer_dim', type=int, default=256, help='number of channels in foreground MLP')
    parser.add_argument('--bg_layer_dim', type=int, default=256, help='number of channels in background MLP')
    parser.add_argument('--appearance_dim', type=int, default=48,
                        help='dimension of appearance embedding vector (set to 0 to disable)')
    parser.add_argument('--affine_appearance', default=False, action='store_true',
                        help='set to true to use affine transformation for appearance instead of latent embedding')

    parser.add_argument('--use_cascade', default=False, action='store_true',
                        help='use separate MLPs to query coarse and fine samples')

    parser.add_argument('--train_mega_nerf', type=str, default=None,
                        help='directory train a Mega-NeRF architecture (point this towards the params.pt file generated by create_cluster_masks.py)')
    parser.add_argument('--boundary_margin', type=float, default=1.15,
                        help='overlap factor between different spatial cells')
    parser.add_argument('--all_val', default=False, action='store_true',
                        help='use all pixels for validation images instead of those specified in cluster masks')
    parser.add_argument('--cluster_2d', default=False, action='store_true', help='cluster without altitude dimension')

    parser.add_argument('--sh_deg', type=int, default=None,
                        help='use spherical harmonics (pos_dir_dim should be set to 0)')

    parser.add_argument('--no_center_pixels', dest='center_pixels', default=True, action='store_false',
                        help='do not shift pixels by +0.5 when computing ray directions')
    parser.add_argument('--no_shifted_softplus', dest='shifted_softplus', default=True, action='store_false',
                        help='use ReLU instead of shifted softplus activation')

    parser.add_argument('--batch_size', type=int, default=1024, help='batch size')
    parser.add_argument('--image_pixel_batch_size', type=int, default=64 * 1024,
                        help='number of pixels to evaluate per split when rendering validation images')
    parser.add_argument('--model_chunk_size', type=int, default=32 * 1024,
                        help='chunk size to split the input to avoid OOM')

    parser.add_argument('--perturb', type=float, default=1.0, help='factor to perturb depth sampling points')
    parser.add_argument('--noise_std', type=float, default=1.0, help='std dev of noise added to regularize sigma')

    parser.add_argument('--lr', type=float, default=5e-4, help='learning rate')
    parser.add_argument('--lr_decay_factor', type=float, default=0.1, help='learning rate decay factor')

    parser.add_argument('--no_bg_nerf', dest='bg_nerf', default=True, action='store_false',
                        help='do not use background MLP')

    parser.add_argument('--ellipse_scale_factor', type=float, default=1.1, help='Factor to scale foreground bounds')
    parser.add_argument('--no_ellipse_bounds', dest='ellipse_bounds', default=True, action='store_false',
                        help='use spherical foreground bounds instead of ellipse')

    parser.add_argument('--train_iterations', type=int, default=500000, help='training iterations')
    parser.add_argument('--val_interval', type=int, default=500001, help='validation interval')
    parser.add_argument('--ckpt_interval', type=int, default=10000, help='checkpoint interval')

    parser.add_argument('--no_resume_ckpt_state', dest='resume_ckpt_state', default=True, action='store_false')

    parser.add_argument('--no_amp', dest='amp', default=True, action='store_false')
    parser.add_argument('--detect_anomalies', default=False, action='store_true')
    parser.add_argument('--random_seed', type=int, default=42)

    return parser
